#!/usr/bin/python3
# coding=utf-8

# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import os
import sys
import logging
import math
from concurrent.futures import ProcessPoolExecutor

import numpy as np

IS_OUTPUT_TXT = False


def compare_chunk(chunk1_data, chunk2_data):
    wrong_num = 0
    eps = 1e-3
    for i in range(chunk1_data.shape[0]):
        a = chunk1_data[i]
        b = chunk2_data[i]
        ae = abs(a - b)
        re = 0 if math.isclose(b, 0) else ae / abs(b)
        if np.isnan(a) or np.isnan(b):
            wrong_num += 1
            continue
        if (ae > eps and re > eps):
            wrong_num += 1
    return wrong_num


def compare_data(work_dir, data_type_str, golden_file="golden.bin", output_file="output.bin"):
    golden_file_path = os.path.join(work_dir, "output", golden_file)
    output_file_path = os.path.join(work_dir, "output", output_file)
    if not os.path.exists(golden_file_path):
        logging.info("[ERROR] can't get golden bin file.")
        return -1
    if not os.path.exists(output_file_path):
        logging.info("[ERROR] can't get output bin file.")
        return -1

    if data_type_str == "float16_float32":
        golden_data = np.fromfile(golden_file_path, dtype="float32")
        output_data = np.fromfile(output_file_path, dtype="float32")
    elif data_type_str == "float16_float16" or data_type_str == "int8_float16_dequant":
        golden_data = np.fromfile(golden_file_path, dtype="float16")
        output_data = np.fromfile(output_file_path, dtype="float16")
    elif data_type_str == "int8_int32_sparse" or data_type_str == "int4_int32":
        golden_data = np.fromfile(golden_file_path, dtype="int32")
        output_data = np.fromfile(output_file_path, dtype="int32")
    else:
        logging.info(f"[ERROR] can't support data type {data_type_str}")
        return -1

    if IS_OUTPUT_TXT:
        np.savetxt(work_dir + "/output/output.txt", output_data.astype(np.float32).flatten(), fmt='%f', newline='\n')

    num_chunks = 32 # process numbers
    total_wrong_num = 0
    chunks1 = np.array_split(output_data, num_chunks)
    chunks2 = np.array_split(golden_data, num_chunks)
    with ProcessPoolExecutor() as executor:
        futures = {executor.submit(compare_chunk, chunk1, chunk2): (chunk1, chunk2)
                   for chunk1, chunk2 in zip(chunks1, chunks2)}
        for future in futures:
            total_wrong_num += future.result()

    return total_wrong_num
